Summary

Pattern match HF style RoPE and delegate to XNNPACK. Relax slice constraints and add unsqueeze, as well, since XNNPACK supports them and it's easiest to not have to conditionally partition them for RoPE only when fusable.

This gives 1-6% performance uplift across decode and prefill for LLMs when fused (HF models, primarily).

Matched pattern:

  def rotate_half(x):
      x1 = x[..., : x.shape[-1] // 2]
      x2 = x[..., x.shape[-1] // 2 :]
      return torch.cat((-x2, x1), dim=-1)

  def apply_rotary_pos_emb(q, k, cos, sin, ...):
      q_embed = (q * cos) + (rotate_half(q) * sin)
      k_embed = (k * cos) + (rotate_half(k) * sin)

See modeling_qwen3.py for an example. I checked 18 LLMs from transformers and about 3/4 use the above pattern.

XNNPACK's RoPE op logically does the following:

def xnnpack_rope(x, weights):
  real = x[..., :d//2]
  imag = x[..., d//2:]
  cos = weights[..., :d//2]
  sin = weights[..., d//2:]

  out_real = real * cos - imag * sin
  out_imag = real * sin + imag * cos
  out = torch.cat([out_real, out_imag], dim=-1)

If the pattern isn't matched, we'll just continue to run decomposed. In order to use the XNNPACK kernel, we do need to combine the cos and sin weights into a single weight tensor. This requires us to modify the weights. Note that the weights depend on position IDs, so this does end up in the graph.

The re-written apply_rotary_pos_emb impl looks like this (post fusion pass). Note that we're assuming that sin and cos second half is duplicated first half. The fusion pass only fuses when it sees that these tensors came from a cat([x, x]) op. We'll conservatively not fuse if this isn't the case. I've tested various HF models and found that it does fuse for the common pattern that most models use.

  def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
      half = cos.shape[-1] // 2
      weights = torch.cat([cos[..., :half], sin[..., :half]], dim=-1)
      weights = weights.unsqueeze(unsqueeze_dim)
      q_embed = xnnpack_rope(q, weights)
      k_embed = xnnpack_rope(k, weights)
      return q_embed, k_embed

In theory, we could do even more aggressive re-writing to fuse even more, but this seems sufficient for now.

cc @digantdesai @cbilgin